# _*_ coding: utf-8 _*_
# @Time : 2021/8/22 00:07
# @Author : xupeng
# contact: ipeng_x1029@163.com
# @File : train.py
import time
from utils import get_time_dif
import numpy as np
import torch
import torch.nn as nn
from sklearn import metrics
import torch.nn.functional as F


# 权重初始化，默认xavier
def init_network(model, method='xavier', exclude='embedding', seed=123):
    for name, w in model.named_parameters():
        if exclude not in name:
            if 'weight' in name:
                if method == 'xavier':
                    nn.init.xavier_normal_(w)
                elif method == 'kaiming':
                    nn.init.kaiming_normal_(w)
                else:
                    nn.init.normal_(w)
            elif 'bias' in name:
                nn.init.constant_(w, 0)
            else:
                pass

def train(config, model, train_iter, test_iter, dev_iter):
    start_time = time.time()
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    loss = nn.CrossEntropyLoss()
    total_batch = 0
    dev_best_loss = float("inf")
    last_improve = 0
    flag = False

    #写入日志：TODO

    for epoch in range(config.num_epochs):
        print(f'Epoch [({epoch+1})/({config.num_epochs})]')
        for i, (trains, label) in enumerate(train_iter):
            optimizer.zero_grad(set_to_none=True)
            outputs = model(trains)
            l = loss(outputs, label)
            l.backward()
            optimizer.step()
            if total_batch % 100 == 0:
                true = label.cpu().detach().numpy()
                pred = outputs.cpu().detach()
                pred = pred.argmax(dim=1).numpy()
                train_acc = metrics.accuracy_score(true, pred)
                dev_acc, dev_loss = evaluate(config, model, dev_iter)
                if dev_loss < dev_best_loss:
                    dev_best_loss = dev_loss
                    torch.save(model.state_dict(), config.save_path)
                    improve = "是"
                    last_improve = total_batch
                else:
                    improve = '否'
                time_gif = get_time_dif(start_time)
                msg = 'Iter:{0}, Train Loss:{1:.2}, Train Acc:{2:.2%}, ' \
                      'Val Loss:{3:.2}, Val Acc:{4:.2%}, Time:{5} 这轮是否更新：{6}'
                print(msg.format(total_batch, l.item(), train_acc, dev_loss, dev_acc, time_gif, improve))
                model.train()
            total_batch += 1
            if total_batch - last_improve > config.nums_need_cut:
                print("No optimization for a long time, early-stopping!")
                flag = True
                break
        if flag:
            break
    test(config, model, test_iter)


def test(config, model, test_iter):
    model.load_state_dict(torch.load(config.save_path))
    model.eval()
    start_time = time.time()
    test_acc, test_loss, test_report, test_confusion = evaluate(config, model, test_iter, is_test=True)
    msg = "Test Loss : {0:.2},     Test Acc:{1:.2%}"
    print(msg.format(test_loss.item(), test_acc))
    print("Precision, Recall and F1-score : ")
    print(test_report)
    print("Confusion Matrix : ")
    print(test_confusion)
    time_gif = get_time_dif(start_time)
    print("Time usage:", time_gif)

def evaluate(config, model, data_iter, is_test=False):
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    with torch.no_grad():
        for datas, labels in data_iter:
            outputs = model(datas)
            loss = F.cross_entropy(outputs, labels)
            loss_total += loss
            labels = labels.cpu().detach().numpy()
            preds = torch.max(outputs.detach(), dim=1)[1].cpu().numpy()
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, preds)
    acc = metrics.accuracy_score(predict_all, labels_all)
    if is_test:
        report = metrics.classification_report(labels_all, predict_all, target_names=config.class_list, digits=4)
        confusion = metrics.classification_report(predict_all, labels_all)
        return acc, loss_total / len(data_iter), report, confusion
    return acc, loss_total / len(data_iter)